package cn.initcap.algorithm.graph.util;

import cn.initcap.algorithm.graph.Edge;
import cn.initcap.algorithm.graph.Graph;
import cn.initcap.algorithm.sort.heap.MinHeap;
import java.util.ArrayList;
import java.util.List;

/**
 * 使用Prim算法求图的最小生成树
 *
 * @author initcap
 * @date Created in 1/20/19 8:53 PM.
 */
public class LazyPrimMst<Weight extends Number & Comparable> {

    /**
     * 图的引用
     */
    private Graph graph;
    /**
     * 最小堆, 算法辅助数据结构
     */
    private MinHeap<Edge<Weight>> pq;
    /**
     * 标记数组, 在算法运行过程中标记节点i是否被访问
     */
    private boolean[] marked;
    /**
     * 最小生成树所包含的所有边
     */
    private List<Edge<Weight>> mst;
    /**
     * 最小生成树的权值
     */
    private Number mstWeight;

    /**
     * 构造函数, 使用Prim算法求图的最小生成树
     *
     * @param graph
     */
    public LazyPrimMst(Graph graph) {

        // 算法初始化
        this.graph = graph;
        pq = new MinHeap<>(this.graph.marginSize());
        marked = new boolean[this.graph.nodeSize()];
        mst = new ArrayList<>();

        // Lazy Prim
        visit(0);
        while (!pq.isEmpty()) {
            // 使用最小堆找出已经访问的边中权值最小的边
            Edge<Weight> e = pq.extractMin();
            // 如果这条边的两端都已经访问过了, 则扔掉这条边
            if (marked[e.v()] == marked[e.w()]) {
                continue;
            }
            // 否则, 这条边则应该存在在最小生成树中
            mst.add(e);

            // 访问和这条边连接的还没有被访问过的节点
            if (!marked[e.v()]) {
                visit(e.v());
            } else {
                visit(e.w());
            }
        }

        // 计算最小生成树的权值
        mstWeight = mst.get(0).wt();
        for (int i = 1; i < mst.size(); i++) {
            mstWeight = mstWeight.doubleValue() + mst.get(i).wt().doubleValue();
        }
    }

    /**
     * 访问节点v
     *
     * @param v
     */
    private void visit(int v) {

        assert !marked[v];
        marked[v] = true;

        // 将和节点v相连接的所有未访问的边放入最小堆中
        for (Object item : graph.adj(v)) {
            Edge<Weight> e = (Edge<Weight>) item;
            if (!marked[e.other(v)]) {
                pq.insert(e);
            }
        }
    }

    /**
     * 返回最小生成树的所有边
     *
     * @return
     */
    public List<Edge<Weight>> mstEdges() {
        return mst;
    }

    /**
     * 返回最小生成树的权值
     *
     * @return
     */
    public Number result() {
        return mstWeight;
    }

}
